"""
knowledge_client.py
知识库客户端封装，基于 ChromaDB 向量数据库实现。
"""

from __future__ import annotations

import logging
import os
import uuid
from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence, Tuple, Union, Any

import chromadb
from chromadb.api import Collection
from chromadb.config import Settings

from .EmbeddingClient import EmbeddingClient
from .RerankClient import RerankClient

logger = logging.getLogger(__name__)


# ------------------ 自定义异常 ------------------
class KnowledgeBaseError(RuntimeError):
    """知识库操作相关的异常基类。"""


class CollectionNotFoundError(KnowledgeBaseError):
    """指定的集合不存在。"""


class DocumentNotFoundError(KnowledgeBaseError):
    """指定的文档不存在。"""


# ------------------ 数据模型 ------------------
@dataclass(slots=True)
class SearchResult:
    """知识库搜索结果，字段清晰可读。"""

    document_id: str
    text: str
    metadata: Dict[str, Any]
    score: float

    def __str__(self) -> str:
        return f"SearchResult(id={self.document_id}, score={self.score:.4f}, metadata={self.metadata})"


# ------------------ 核心封装 ------------------
class KnowledgeClient:
    """
    知识库客户端，基于 ChromaDB 向量数据库实现。
    
    提供知识库的创建、文档管理、相似度搜索等功能，支持元数据过滤和结果重排序。
    
    用法示例
    --------
    >>> # 创建知识库客户端
    >>> client = KnowledgeClient(
    ...     persist_directory="./db",
    ...     embedding_client=EmbeddingClient()
    ... )
    >>> 
    >>> # 创建知识库集合
    >>> client.create_collection("my_docs")
    >>> 
    >>> # 添加文档
    >>> client.add_documents(
    ...     collection_name="my_docs",
    ...     texts=["这是第一篇文档", "这是第二篇文档"],
    ...     metadatas=[{"source": "book1"}, {"source": "book2"}],
    ...     ids=["doc1", "doc2"]
    ... )
    >>> 
    >>> # 搜索文档
    >>> results = client.search(
    ...     collection_name="my_docs",
    ...     query="文档",
    ...     n_results=5,
    ...     where={"source": "book1"}
    ... )
    >>> for result in results:
    ...     print(f"{result.document_id}: {result.text} (score: {result.score})")
    """

    def __init__(
        self,
        persist_directory: str = "./db",
        embedding_client: Optional[EmbeddingClient] = None,
        rerank_client: Optional[RerankClient] = None,
        *,
        collection_metadata: Optional[Dict[str, Any]] = None,
    ):
        """
        初始化知识库客户端。
        
        Parameters
        ----------
        persist_directory : str, default="./db"
            ChromaDB 数据持久化目录路径。
        embedding_client : Optional[EmbeddingClient], default=None
            用于生成文本嵌入的客户端。如果为 None，将创建默认的 EmbeddingClient。
        rerank_client : Optional[RerankClient], default=None
            用于结果重排序的客户端。如果为 None，将不进行重排序。
        collection_metadata : Optional[Dict[str, Any]], default=None
            集合元数据，用于存储集合的附加信息。
        """
        # 初始化 ChromaDB 客户端
        self._chroma_client = chromadb.PersistentClient(
            path=persist_directory,
            settings=Settings(
                anonymized_telemetry=False,
                allow_reset=True,
            )
        )
        
        # 初始化嵌入客户端
        self._embedding_client = embedding_client or EmbeddingClient()
        
        # 初始化重排序客户端（可选）
        self._rerank_client = rerank_client
        
        # 集合元数据
        self._collection_metadata = collection_metadata or {}
        
        logger.debug("KnowledgeClient initialized with persist_directory=%s", persist_directory)

    def list_collections(self) -> List[str]:
        """
        列出所有知识库集合名称。
        
        Returns
        -------
        List[str]
            集合名称列表。
        """
        collections = self._chroma_client.list_collections()
        return [collection.name for collection in collections]

    def create_collection(
        self,
        collection_name: str,
        metadata: Optional[Dict[str, Any]] = None,
        get_or_create: bool = True,
    ) -> Collection:
        """
        创建新的知识库集合。
        
        Parameters
        ----------
        collection_name : str
            集合名称，必须唯一。
        metadata : Optional[Dict[str, Any]], default=None
            集合元数据，用于存储集合的附加信息。
        get_or_create : bool, default=True
            如果为 True，当集合已存在时返回现有集合；如果为 False，当集合已存在时抛出异常。
        
        Returns
        -------
        Collection
            创建的 ChromaDB 集合对象。
        
        Raises
        ------
        ValueError
            如果 get_or_create=False 且集合已存在。
        """
        # 合并默认元数据和自定义元数据
        combined_metadata = {**self._collection_metadata}
        if metadata:
            combined_metadata.update(metadata)
            
        try:
            collection = self._chroma_client.create_collection(
                name=collection_name,
                metadata=combined_metadata,
                embedding_function=self._get_embedding_function(),
            )
            logger.info("Created collection: %s", collection_name)
            return collection
        except ValueError as e:
            if "already exists" in str(e) and get_or_create:
                logger.debug("Collection %s already exists, getting existing collection", collection_name)
                return self._chroma_client.get_collection(
                    name=collection_name,
                    embedding_function=self._get_embedding_function(),
                )
            raise

    def get_collection(self, collection_name: str) -> Collection:
        """
        获取现有的知识库集合。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        
        Returns
        -------
        Collection
            ChromaDB 集合对象。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        """
        try:
            return self._chroma_client.get_collection(
                name=collection_name,
                embedding_function=self._get_embedding_function(),
            )
        except ValueError as e:
            if "does not exist" in str(e):
                raise CollectionNotFoundError(f"Collection '{collection_name}' does not exist") from e
            raise

    def delete_collection(self, collection_name: str) -> None:
        """
        删除知识库集合。
        
        Parameters
        ----------
        collection_name : str
            要删除的集合名称。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        """
        try:
            self._chroma_client.delete_collection(name=collection_name)
            logger.info("Deleted collection: %s", collection_name)
        except ValueError as e:
            if "does not exist" in str(e):
                raise CollectionNotFoundError(f"Collection '{collection_name}' does not exist") from e
            raise

    def add_documents(
        self,
        collection_name: str,
        texts: List[str],
        metadatas: Optional[List[Dict[str, Any]]] = None,
        ids: Optional[List[str]] = None,
        *,
        chunk_size: Optional[int] = None,
        chunk_overlap: int = 0,
    ) -> List[str]:
        """
        向知识库集合添加文档。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        texts : List[str]
            要添加的文本内容列表。
        metadatas : Optional[List[Dict[str, Any]]], default=None
            文档元数据列表，与 texts 一一对应。如果为 None，则使用空字典。
        ids : Optional[List[str]], default=None
            文档 ID 列表，与 texts 一一对应。如果为 None，则自动生成 UUID。
        chunk_size : Optional[int], default=None
            文本分块大小（字符数）。如果指定，将对长文本进行分块。
        chunk_overlap : int, default=0
            分块重叠大小（字符数）。
        
        Returns
        -------
        List[str]
            添加的文档 ID 列表。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        ValueError
            如果 texts、metadatas 和 ids 的长度不一致。
        """
        # 获取集合
        collection = self.get_collection(collection_name)
        
        # 处理文档 ID
        if ids is None:
            ids = [str(uuid.uuid4()) for _ in range(len(texts))]
        elif len(ids) != len(texts):
            raise ValueError("Length of ids must match length of texts")
        
        # 处理元数据
        if metadatas is None:
            metadatas = [{} for _ in range(len(texts))]
        elif len(metadatas) != len(texts):
            raise ValueError("Length of metadatas must match length of texts")
        
        # 如果需要分块
        if chunk_size is not None and chunk_size > 0:
            chunked_texts = []
            chunked_metadatas = []
            chunked_ids = []
            
            for i, (text, metadata, doc_id) in enumerate(zip(texts, metadatas, ids)):
                # 简单的字符分块，实际应用中可能需要更复杂的分块策略
                if len(text) <= chunk_size:
                    chunked_texts.append(text)
                    chunked_metadatas.append(metadata)
                    chunked_ids.append(doc_id)
                else:
                    chunks = []
                    for j in range(0, len(text), chunk_size - chunk_overlap):
                        chunk = text[j:j + chunk_size]
                        if len(chunk) < 10:  # 忽略过小的块
                            continue
                        chunks.append(chunk)
                        
                        # 为每个块创建元数据和 ID
                        chunk_metadata = metadata.copy()
                        chunk_metadata["chunk_index"] = len(chunks) - 1
                        chunk_metadata["original_id"] = doc_id
                        chunked_metadatas.append(chunk_metadata)
                        
                        chunk_id = f"{doc_id}_chunk_{len(chunks) - 1}"
                        chunked_ids.append(chunk_id)
                    
                    chunked_texts.extend(chunks)
            
            texts = chunked_texts
            metadatas = chunked_metadatas
            ids = chunked_ids
        
        # 添加文档
        collection.add(
            documents=texts,
            metadatas=metadatas,
            ids=ids,
        )
        
        logger.info("Added %d documents to collection %s", len(texts), collection_name)
        return ids

    def get_document(
        self,
        collection_name: str,
        document_id: str,
    ) -> Tuple[str, Dict[str, Any]]:
        """
        获取知识库中的单个文档。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        document_id : str
            文档 ID。
        
        Returns
        -------
        Tuple[str, Dict[str, Any]]
            文档内容和元数据的元组。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        DocumentNotFoundError
            如果指定的文档不存在。
        """
        collection = self.get_collection(collection_name)
        
        result = collection.get(
            ids=[document_id],
            include=["documents", "metadatas"],
        )
        
        if not result["ids"]:
            raise DocumentNotFoundError(f"Document '{document_id}' not found in collection '{collection_name}'")
        
        return result["documents"][0], result["metadatas"][0]

    def delete_document(
        self,
        collection_name: str,
        document_id: str,
    ) -> None:
        """
        从知识库中删除单个文档。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        document_id : str
            文档 ID。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        """
        collection = self.get_collection(collection_name)
        collection.delete(ids=[document_id])
        logger.info("Deleted document %s from collection %s", document_id, collection_name)

    def update_document(
        self,
        collection_name: str,
        document_id: str,
        text: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        更新知识库中的文档。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        document_id : str
            文档 ID。
        text : str
            新的文档内容。
        metadata : Optional[Dict[str, Any]], default=None
            新的文档元数据。如果为 None，则保留原有元数据。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        DocumentNotFoundError
            如果指定的文档不存在。
        """
        collection = self.get_collection(collection_name)
        
        # 如果需要保留原有元数据
        if metadata is None:
            try:
                _, original_metadata = self.get_document(collection_name, document_id)
                metadata = original_metadata
            except DocumentNotFoundError:
                metadata = {}
        
        # 更新文档
        collection.update(
            ids=[document_id],
            documents=[text],
            metadatas=[metadata],
        )
        
        logger.info("Updated document %s in collection %s", document_id, collection_name)

    def search(
        self,
        collection_name: str,
        query: str,
        n_results: int = 5,
        where: Optional[Dict[str, Any]] = None,
        where_document: Optional[Dict[str, Any]] = None,
        *,
        rerank: bool = False,
    ) -> List[SearchResult]:
        """
        在知识库中搜索与查询相似的文档。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        query : str
            查询文本。
        n_results : int, default=5
            返回的结果数量。
        where : Optional[Dict[str, Any]], default=None
            元数据过滤条件，用于筛选特定元数据的文档。
        where_document : Optional[Dict[str, Any]], default=None
            文档内容过滤条件，用于筛选特定内容的文档。
        rerank : bool, default=False
            是否使用重排序客户端对结果进行重排序。
            如果为 True 且 self._rerank_client 不为 None，则进行重排序。
        
        Returns
        -------
        List[SearchResult]
            搜索结果列表，按相似度降序排序。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        ValueError
            如果 rerank=True 但未提供 rerank_client。
        """
        collection = self.get_collection(collection_name)
        
        # 执行向量搜索
        results = collection.query(
            query_texts=[query],
            n_results=n_results,
            where=where,
            where_document=where_document,
            include=["documents", "metadatas", "distances"],
        )
        
        # 处理搜索结果
        search_results = []
        if results["ids"] and results["ids"][0]:
            for i, (doc_id, text, metadata, distance) in enumerate(zip(
                results["ids"][0],
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0],
            )):
                # 将距离转换为相似度分数（1 - 归一化距离）
                # ChromaDB 使用余弦距离，范围为 0-2，0 表示完全相同
                score = 1.0 - (distance / 2.0)
                
                search_results.append(
                    SearchResult(
                        document_id=doc_id,
                        text=text,
                        metadata=metadata,
                        score=score,
                    )
                )
        
        # 如果需要重排序
        if rerank and self._rerank_client is not None:
            if not search_results:
                return search_results
                
            # 提取文档文本
            texts = [result.text for result in search_results]
            
            # 使用重排序客户端
            reranked = self._rerank_client.rerank(query, texts)
            
            # 更新分数并重新排序
            for i, reranked_result in enumerate(reranked):
                search_results[reranked_result.index].score = reranked_result.score
            
            # 按新分数排序
            search_results.sort(key=lambda x: x.score, reverse=True)
        elif rerank and self._rerank_client is None:
            raise ValueError("rerank=True but no rerank_client provided during initialization")
        
        logger.debug("Found %d results for query '%s' in collection %s", len(search_results), query, collection_name)
        return search_results

    def search_by_vector(
        self,
        collection_name: str,
        embedding: List[float],
        n_results: int = 5,
        where: Optional[Dict[str, Any]] = None,
        where_document: Optional[Dict[str, Any]] = None,
    ) -> List[SearchResult]:
        """
        使用向量直接在知识库中搜索相似文档。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        embedding : List[float]
            查询向量。
        n_results : int, default=5
            返回的结果数量。
        where : Optional[Dict[str, Any]], default=None
            元数据过滤条件，用于筛选特定元数据的文档。
        where_document : Optional[Dict[str, Any]], default=None
            文档内容过滤条件，用于筛选特定内容的文档。
        
        Returns
        -------
        List[SearchResult]
            搜索结果列表，按相似度降序排序。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        """
        collection = self.get_collection(collection_name)
        
        # 执行向量搜索
        results = collection.query(
            query_embeddings=[embedding],
            n_results=n_results,
            where=where,
            where_document=where_document,
            include=["documents", "metadatas", "distances"],
        )
        
        # 处理搜索结果
        search_results = []
        if results["ids"] and results["ids"][0]:
            for i, (doc_id, text, metadata, distance) in enumerate(zip(
                results["ids"][0],
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0],
            )):
                # 将距离转换为相似度分数（1 - 归一化距离）
                score = 1.0 - (distance / 2.0)
                
                search_results.append(
                    SearchResult(
                        document_id=doc_id,
                        text=text,
                        metadata=metadata,
                        score=score,
                    )
                )
        
        logger.debug("Found %d results for vector search in collection %s", len(search_results), collection_name)
        return search_results

    def count_documents(self, collection_name: str) -> int:
        """
        获取知识库集合中的文档数量。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        
        Returns
        -------
        int
            文档数量。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        """
        collection = self.get_collection(collection_name)
        return collection.count()

    def get_collection_metadata(self, collection_name: str) -> Dict[str, Any]:
        """
        获取知识库集合的元数据。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        
        Returns
        -------
        Dict[str, Any]
            集合元数据。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        """
        collection = self.get_collection(collection_name)
        return collection.metadata or {}

    def update_collection_metadata(
        self,
        collection_name: str,
        metadata: Dict[str, Any],
    ) -> None:
        """
        更新知识库集合的元数据。
        
        Parameters
        ----------
        collection_name : str
            集合名称。
        metadata : Dict[str, Any]
            新的集合元数据。
        
        Raises
        ------
        CollectionNotFoundError
            如果指定的集合不存在。
        """
        collection = self.get_collection(collection_name)
        collection.modify(metadata=metadata)
        logger.info("Updated metadata for collection %s", collection_name)

    def _get_embedding_function(self):
        """
        创建一个适用于 ChromaDB 的嵌入函数。
        
        Returns
        -------
        EmbeddingFunction
            符合 ChromaDB 接口的嵌入函数对象。
        """
        # 创建一个符合 ChromaDB EmbeddingFunction 接口的类
        class EmbeddingFunction:
            def __init__(self, embedding_client):
                self._embedding_client = embedding_client
                
            def __call__(self, input: List[str]) -> List[List[float]]:
                """
                ChromaDB 要求的嵌入函数接口。
                
                Parameters
                ----------
                input : List[str]
                    要嵌入的文本列表。
                
                Returns
                -------
                List[List[float]]
                    嵌入向量列表。
                """
                return self._embedding_client.embed_batch(input)
                
            def name(self) -> str:
                """
                返回嵌入函数的名称，用于 ChromaDB 识别。
                
                Returns
                -------
                str
                    嵌入函数的名称。
                """
                return f"custom_embedding_function_{self._embedding_client._model}"
        
        return EmbeddingFunction(self._embedding_client)